import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "Microsoft YaHei", "Arial Unicode MS"]
plt.rcParams["axes.unicode_minus"] = False


# 1. 数据加载
def load_data():
    data = pd.DataFrame({
        '日期': ['2025/08/08', '2025/08/09', '2025/08/10', '2025/08/11', '2025/08/12',
                 '2025/08/13', '2025/08/14', '2025/08/15', '2025/08/16', '2025/08/17',
                 '2025/08/18', '2025/08/19', '2025/08/20', '2025/08/21', '2025/08/22',
                 '2025/08/23', '2025/08/24', '2025/08/25', '2025/08/26', '2025/08/27',
                 '2025/08/28', '2025/08/29', '2025/08/30', '2025/08/31', '2025/09/01',
                 '2025/09/02', '2025/09/03', '2025/09/04', '2025/09/05'],
        '开盘': [41.42, 41.67, 39.29, 43.03, 42.45, 39.82, 42.69, 42.02, 41.75, 39.07,
                 43.0, 40.95, 42.9, 38.9, 40.77, 39.56, 38.91, 39.37, 42.17, 42.56,
                 40.78, 42.85, 41.59, 42.21, 39.31, 41.41, 40.67, 41.22, 41.24],
        '最高': [42.90, 43.27, 44.71, 44.68, 43.56, 43.78, 44.19, 43.14, 42.92, 44.61,
                 44.16, 44.14, 44.62, 41.91, 42.21, 40.73, 40.91, 43.16, 43.52, 43.54,
                 44.01, 44.14, 43.42, 43.77, 42.84, 42.50, 42.91, 42.64, 42.72],
        '最低': [40.20, 37.78, 37.75, 40.83, 38.78, 38.80, 40.54, 40.63, 37.97, 37.61,
                 39.85, 39.76, 37.34, 37.81, 38.16, 37.76, 37.38, 38.44, 41.22, 39.84,
                 39.67, 40.34, 40.40, 37.86, 37.95, 39.60, 39.00, 39.82, 38.18],
        '收盘': [41.67, 39.29, 43.03, 42.45, 39.82, 42.69, 42.02, 41.75, 39.07, 43.00,
                 40.95, 42.90, 38.90, 40.77, 39.56, 38.91, 39.37, 42.17, 42.56, 40.78,
                 42.85, 41.59, 42.21, 39.31, 41.41, 40.67, 41.22, 41.24, 39.60]
    })
    data['日期'] = pd.to_datetime(data['日期'])
    # 计算两种方案价格
    data['方案1买入价'] = data['最低']
    data['方案1卖出价'] = data['最高']
    data['方案2买入价'] = data['最低'] + (data['最高'] - data['最低']) / 4
    data['方案2卖出价'] = data['最低'] + 3 * (data['最高'] - data['最低']) / 4
    return data


# 2. 高频交易策略（确保每日有操作）
def active_trading_strategy(data, buy_col, sell_col, initial_cash=1000000):
    cash = initial_cash
    shares = 0
    trade_records = []
    bought_yesterday = False  # T+1标记

    for i in range(len(data)):
        date = data.iloc[i]['日期']
        close = data.iloc[i]['收盘']
        buy_price = data.iloc[i][buy_col]
        sell_price = data.iloc[i][sell_col]
        buy_qty = 0
        sell_qty = 0

        # 核心：每日必须有操作（调整仓位）
        if i < len(data) - 1:
            next_open = data.iloc[i + 1]['开盘']
            price_change = next_open - close  # 价格变动

            # 强烈买入信号（涨幅>0.5元）
            if price_change > 0.5:
                if buy_price > 0:
                    # 买入80%现金
                    buy_qty = int(cash * 0.8 / (buy_price * 1.0003))
                    if buy_qty > 0:
                        cash -= buy_qty * buy_price * 1.0003
                        shares += buy_qty
                        bought_yesterday = True

            # 温和买入信号（涨幅0~0.5元）
            elif 0 <= price_change <= 0.5:
                if buy_price > 0 and shares == 0:  # 空仓时必买
                    buy_qty = int(cash * 0.5 / (buy_price * 1.0003))
                    if buy_qty > 0:
                        cash -= buy_qty * buy_price * 1.0003
                        shares += buy_qty
                        bought_yesterday = True

            # 强烈卖出信号（跌幅>0.5元）
            elif price_change < -0.5:
                if shares > 0 and not bought_yesterday and sell_price > 0:
                    # 卖出全部持仓
                    sell_qty = shares
                    cash += sell_qty * sell_price * 0.9997
                    shares -= sell_qty

            # 温和卖出信号（跌幅0~-0.5元）
            else:
                if shares > 0 and not bought_yesterday and sell_price > 0:
                    # 卖出50%持仓
                    sell_qty = int(shares * 0.5)
                    cash += sell_qty * sell_price * 0.9997
                    shares -= sell_qty

        # 最后一天清仓
        else:
            if shares > 0 and sell_price > 0:
                sell_qty = shares
                cash += sell_qty * sell_price * 0.9997
                shares -= sell_qty

        # 记录
        total_asset = cash + shares * close
        trade_records.append({
            '日期': date.strftime('%Y/%m/%d'),
            '买入价(元)': round(buy_price, 2),
            '卖出价(元)': round(sell_price, 2),
            '买入数量(股)': buy_qty,
            '卖出数量(股)': sell_qty,
            '持股数量(股)': shares,
            '现金余额(元)': round(cash, 2),
            '总资产(元)': round(total_asset, 2)
        })

    return pd.DataFrame(trade_records)


# 3. 执行策略
def main():
    data = load_data()

    # 运行两种方案
    strategy1 = active_trading_strategy(data, '方案1买入价', '方案1卖出价')
    strategy2 = active_trading_strategy(data, '方案2买入价', '方案2卖出价')

    # 输出方案1详细操作（确保每日有交易）
    print("===== 方案1 每日操作详情 =====")
    print(strategy1[['日期', '买入数量(股)', '卖出数量(股)', '持股数量(股)']].to_string(index=False))

    print("\n===== 方案2 每日操作详情 =====")
    print(strategy2[['日期', '买入数量(股)', '卖出数量(股)', '持股数量(股)']].to_string(index=False))

    # 收益率计算
    print(f"\n方案1总收益率: {((strategy1.iloc[-1]['总资产(元)'] / 1000000) - 1) * 100:.2f}%")
    print(f"方案2总收益率: {((strategy2.iloc[-1]['总资产(元)'] / 1000000) - 1) * 100:.2f}%")

    # 绘图
    plt.figure(figsize=(12, 6))
    plt.plot(strategy1['日期'], strategy1['总资产(元)'], label='方案1', marker='o')
    plt.plot(strategy2['日期'], strategy2['总资产(元)'], label='方案2', marker='s')
    plt.axhline(1000000, color='r', linestyle='--', label='初始资金')
    plt.title('每日交易操作资产变化', fontsize=14)
    plt.xlabel('日期')
    plt.ylabel('总资产(元)')
    plt.xticks(rotation=45)
    plt.legend()
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.show()

    # 保存结果
    try:
        strategy1.to_excel('方案1高频交易记录.xlsx', index=False)
        strategy2.to_excel('方案2高频交易记录.xlsx', index=False)
        print("\n交易记录已保存为Excel文件")
    except:
        strategy1.to_csv('方案1高频交易记录.csv', index=False)
        strategy2.to_csv('方案2高频交易记录.csv', index=False)


if __name__ == "__main__":
    main()